#include <iostream>
#include <string>
#include <cstdlib>
#include <cstring>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

using namespace std;

static const int retry_count = 5;

void Usage(std::string proc)
{
    std::cout << "Usage:\n\t" << proc << " serverip serverport\n"
              << std::endl;
}

bool visitServer(const std::string& serverip, uint16_t serverport, int* retry_count)
{
    // 1. create socket
    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd < 0)
    {
        std::cerr << "Error: Failed to create socket." << std::endl;
        close(sockfd);
        return false;
    }

    // 2. connect to server
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(serverport);
    if (inet_pton(AF_INET, serverip.c_str(), &server_addr.sin_addr) == 0)
    {
        std::cerr << "Error: Invalid server IP address." << std::endl;
        close(sockfd);
        return false;
    }

    int n = connect(sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr));
    if (n < 0)
    {
        std::cerr << "Error: Failed to connect to server." << std::endl;
        close(sockfd);
        return false;
    }

    *retry_count = 0;  // reset retry count

    // 3. send message to server
    while(true)
    {
        string msg;
        cout << "Enter message to send to server (empty line to exit): ";
        getline(cin, msg);
        if (msg.empty())
        {
            close(sockfd);
            return true;
        }
        ssize_t n = write(sockfd, msg.c_str(), msg.size());
        cout << "write success, n = " << n << endl;
        if (n < 0)
        {
            std::cerr << "Error: Failed to send message to server." << std::endl;
            close(sockfd);
            return false;
        }
        else
        {
            string recv_msg(1024, '\0');
            n = read(sockfd, &recv_msg[0], recv_msg.size());
            if (n < 0)
            {
                std::cerr << "Error: Failed to receive message from server." << std::endl;
                close(sockfd);
                return false;
            }
            else if (n == 0)
            {
                close(sockfd);
                return false;
            }
            cout << "Received message from server: " << recv_msg << endl;
        }
    }

    // 4. close socket
    close(sockfd);
    return true;
}

// ./udpclient serverip serverport
int main(int argc, char *argv[])
{
    if (argc != 3)
    {
        Usage(argv[0]);
        exit(1);
    }
    std::string serverip = argv[1];
    uint16_t serverport = std::stoi(argv[2]);

    for(int i = 0; i < retry_count; ++i)
    {
        if (visitServer(serverip, serverport, &i))  // success
        {
            break;
        }
        else
        {
            std::cerr << "Error: Failed to connect to server. Retry count:" << i << std::endl;
            sleep(1);  // wait for 1 second before retrying
        }
    }

    return 0;
}